use std::io;

use thiserror::Error;


#[derive(Debug, Error)]
enum TestError {
    #[error("env error: {0}")]
    VarError(String),

    #[error("HTML to Markdown conversion error: {0}")]
    IoError(#[from] io::Error),

    #[error("Prompt error: {0}")]
    PromptError(#[from] rig::completion::PromptError),
}



#[cfg(test)]
mod tests {
    use std::env;

    use base64::{prelude::BASE64_STANDARD, Engine};
    use rig::{agent::AgentBuilder, completion::Prompt, message::{ContentFormat, Image, ImageDetail, ImageMediaType, Message, Text, UserContent}, providers::gemini, OneOrMany};

    use super::*;

    #[tokio::test]
    async fn test_gemini_img() -> Result<(), TestError> {
        dotenvy::from_path("../../.env").ok();

        let target_file_path = "resources/use.md";
        let content = std::fs::read_to_string(target_file_path)?;


        let img_path = "resources/wechat_2025-04-18_212136_576.png";
        let image_base64 = std::fs::read(img_path).map(|bytes| {
            let encoded = BASE64_STANDARD.encode(bytes);
            encoded
            //format!("data:image/png;base64,{}", encoded)
        })?;

        let image = Image {
            data: image_base64,
            media_type: Some(ImageMediaType::PNG),
            format: Some(ContentFormat::Base64),
            detail: Some(ImageDetail::High),
        };


        let api_key = env::var("GEMINI_API_KEY").map_err(|e| TestError::VarError(e.to_string()))?;
        let base_url = env::var("GEMINI_BASE_URL").unwrap_or_else(|_| "https://generativelanguage.googleapis.com".to_string());
        let client = gemini::Client::from_url(&api_key, &base_url);
        let model = client.completion_model("gemini-2.0-flash");
        let system = r#"
请扮演一个**极其注重细节**的内容分析和格式转换专家。你的任务是结合图像识别和给定的文本上下文信息，生成一份**内容详尽、细节丰富**的结构化 Markdown 文档。

请遵循以下步骤，并**力求在每个步骤中都做到最充分的挖掘和描述**:

1.  **视觉内容提取**: 运用你的视觉识别能力，**极其仔细地**分析我提供的图片。**不仅要**准确、完整地提取图片中的所有可见文本、关键元素和整体布局结构，**还要尽可能详细地描述**这些元素的特征、相互关系以及给你的整体印象。**不要遗漏任何视觉细节**。
2.  **上下文信息整合**: **深入**阅读我提供的上下文（Context）文本。识别并提取其中**所有**与图片内容相关的补充信息，特别是那些图片本身无法直接展示的信息，例如：
    *   图片或媒体资源的具体 URL
    *   相关的超链接地址
    *   **详细的**文字注释或背景说明
    *   上下文中提及但图片未明确显示的**所有**细节
    *   **任何可能有助于理解图片背景或深层含义的信息**
3.  **Markdown 文档生成**: 将上述两部分信息整合，并按照以下要求生成 Markdown 文档：
    *   **主标题**: 使用一个 H1 级别的标题 (`# 标题`) 作为文档的最高层级标题。
    *   **层级结构**: 根据内容的逻辑关系，合理使用 H2 (`##`)、H3 (`###`) 等次级标题来组织内容，形成清晰的层次结构。**在每个层级下，都应进行充分的阐述**。
    *   **内容呈现**: 使用列表（有序或无序）、引用、代码块、**以及必要的描述性段落**来恰当地展示提取和整合后的信息。**对于关键信息点，请展开详细说明，避免过于简略**。
    *   **准确链接**: 确保所有从上下文中补充的 URL 和超链接都准确无误地嵌入到 Markdown 文档中。
    *   **完整全面与详尽**: 确保最终文档不仅准确反映了图片内容、包含了上下文提供的所有相关补充信息，**而且在描述上尽可能地详细和深入**。**目标是生成一份信息量大、内容饱满的文档，长度应足以覆盖所有识别到的细节**。

请确保最终输出的 Markdown 文档内容准确、结构清晰、格式规范，并且**在详细程度上达到最高标准**。
"#;
        let agent = AgentBuilder::new(model)
            .preamble(system)
            .temperature(0.4)
            //.context(&content)
            .build();


        let user_content1 = UserContent::Image(image);
        let user_content2 = UserContent::Text(Text { text: content });

        // 使用OneOrMany::many创建包含多个内容的消息
        let message = Message::User {
            content: OneOrMany::many(vec![user_content1]).unwrap(),
        };

        let response = agent.prompt(message).await?;

        let output_path = "resources/gemini_result.md";
        std::fs::write(output_path, response)?;
        Ok(())
    }

    #[tokio::test]
    async fn test_gemini_text() -> Result<(), TestError> {
        dotenvy::from_path("./.env").ok();

        let api_key = env::var("GEMINI_API_KEY").map_err(|e| TestError::VarError(e.to_string()))?;
        let client = gemini::Client::new(&api_key);
        let model = client.completion_model("gemini-2.5-flash");

        let agent = AgentBuilder::new(model)
            .preamble("你是一个专业的技术顾问，擅长解释技术概念和提供详细的技术建议。")
            .temperature(0.7)
            .build();

        let response = agent.prompt("请解释什么是向量数据库，以及它在AI应用中的作用").await?;

        let output_path = "resources/gemini_text_result.md";
        std::fs::write(output_path, response)?;
        Ok(())
    }
}
